"""A proxy server to help debug/diagnose xapian communications.

This proxy server sits between the xapian server and the client.  This
proxy is started with the remote host and port running the xapian server,
and the local port name that will be listened on for connections.

The client is then instructed to connect to the proxy port, not the
remote server port.  The proxy will then log all communication between
the client and the server.
"""
import sys
import traceback
import re
import datetime
import optparse
import asynchat
import asyncore
import socket

# Pasted from xapian's remoteprotocol.h
message_types_desc = """
    MSG_ALLTERMS,		// All Terms
    MSG_COLLFREQ,		// Get Collection Frequency
    MSG_DOCUMENT,		// Get Document
    MSG_TERMEXISTS,		// Term Exists?
    MSG_TERMFREQ,		// Get Term Frequency
    MSG_KEEPALIVE,		// Keep-alive
    MSG_DOCLENGTH,		// Get Doc Length
    MSG_QUERY,			// Run Query
    MSG_TERMLIST,		// Get TermList
    MSG_POSITIONLIST,		// Get PositionList
    MSG_POSTLIST,		// Get PostList
    MSG_REOPEN,			// Reopen
    MSG_UPDATE,			// Get Updated DocCount and AvLength
    MSG_ADDDOCUMENT,		// Add Document
    MSG_CANCEL,			// Cancel
    MSG_DELETEDOCUMENT,		// Delete Document
    MSG_DELETEDOCUMENTTERM,	// Delete Document by term
    MSG_FLUSH,			// Flush
    MSG_REPLACEDOCUMENT,	// Replace Document
    MSG_REPLACEDOCUMENTTERM,	// Replace Document by term
    MSG_GETMSET,		// Get MSet
    MSG_SHUTDOWN,		// Shutdown
"""
reply_types_desc = """
    REPLY_GREETING,		// Greeting
    REPLY_EXCEPTION,		// Exception
    REPLY_DONE,			// Done sending list
    REPLY_ALLTERMS,		// All Terms
    REPLY_COLLFREQ,		// Get Collection Frequency
    REPLY_DOCDATA,		// Get Document
    REPLY_TERMDOESNTEXIST,	// Term Doesn't Exist
    REPLY_TERMEXISTS,		// Term Exists
    REPLY_TERMFREQ,		// Get Term Frequency
    REPLY_DOCLENGTH,		// Get Doc Length
    REPLY_RESULTS,		// Results (MSet)
    REPLY_STATS,		// Stats
    REPLY_TERMLIST,		// Get Termlist
    REPLY_POSITIONLIST,		// Get PositionList
    REPLY_POSTLISTSTART,	// Start of a postlist
    REPLY_POSTLISTITEM,		// Item in body of a postlist
    REPLY_UPDATE,		// Get Updated DocCount and AvLength
    REPLY_VALUE,		// Document Value
    REPLY_ADDDOCUMENT,		// Add Document
"""
# Build up a list - just take the name
message_types = []
reply_types = []
for (target, descs) in [(message_types, message_types_desc),
                        (reply_types, reply_types_desc)]:
    for mt in descs.splitlines():
        if mt:
            target.append(re.match(r'\W*(\w*),', mt).group(1))
del target, descs, mt

# A mixin class where xapian semantics can go.
class xapian_message_pipe:
    def __init__(self):
        self.buffer = ''
        self.set_terminator(None)
        self.msg_type = None
        self.header_len = 2 # may grow for large transfers.
        self.body_len = None # will get set once.
        self.streaming = False

    def handle_incoming_data(self, data, pusher):
        if self.streaming:
            pusher.push(data)
            self.bytes_pushd += len(data)
            if self.bytes_pushd == self.body_len:
                self.streaming = False
            return None, None

        # not streaming - must still be working out sizes etc.
        self.buffer += data
        if len(self.buffer) < 2:
            return None, None
        self.msg_type = ord(self.buffer[0])
        self.body_len = ord(self.buffer[1])
        self.header_len = 2
        if self.body_len == 0xff:
            if len(self.buffer) < 6:
                return None, None
            # turn it into header and body len
            nbytes = i = shift = 0
            for i in range(4):
                ch = ord(self.buffer[i+2])
                nbytes |= (ch & 0x7f) << shift
                shift += 7
                self.header_len += 1
                if ch & 0x80:
                    break
            else:
                raise RuntimeError, "Insane length"
            self.body_len = nbytes + 255
        pusher.push(self.buffer)
        self.bytes_pushd = len(self.buffer) - self.header_len
        self.buffer = ''
        self.streaming = True
        return self.msg_type, self.header_len + self.body_len

    def handle_error(self):
        print "Error proxying request:"
        traceback.print_exc()
        raise asyncore.ExitNow

    # Use this to print tracing information - it includes a timestamp
    # 'IN' means 'from client, to server'
    IN='--> '
    OUT='<-- '
    def trace(self, direction, msg, *args):
        msg = msg % args
        now = datetime.datetime.now()
        ts = "%s.%03d" % (now.strftime("%X"), now.microsecond / 1000)
        full_msg = "%s[%s, %s] %s\n" % (direction, self.id, ts, msg)
        sys.stdout.write(full_msg)

# proxy code based on asynchat docs at www.nightmare.com
class proxy_server (asyncore.dispatcher):
    def __init__ (self, host, port, proxy_port, options):
        self.options = options
        asyncore.dispatcher.__init__ (self)
        self.create_socket (socket.AF_INET, socket.SOCK_STREAM)
        self.set_reuse_addr()
        self.there = (host, port)
        here = ('', proxy_port)
        self.bind (here)
        self.listen (5)

    def handle_accept (self):
        conn, addr = self.accept()
        proxy_receiver (self, (conn, addr))

    def handle_error(self):
        print "Error proxying request (from server):"
        traceback.print_exc()
        raise asyncore.ExitNow
    
class proxy_sender(xapian_message_pipe, asynchat.async_chat):
    ac_in_buffer_size = 32768
    ac_out_buffer_size = 32768
    def __init__(self, receiver, address, options):
        self.options = options
        asynchat.async_chat.__init__(self)
        xapian_message_pipe.__init__(self)
        self.receiver = receiver
        self.create_socket (socket.AF_INET, socket.SOCK_STREAM)
        self.connect (address)

    def handle_connect (self):
        pass

    def collect_incoming_data (self, data):
        msg_type, nbytes = self.handle_incoming_data(data, self.receiver)
        if msg_type is None:
            return
        self.trace(self.OUT, "%s (%d bytes)", reply_types[msg_type], nbytes)

    def handle_close (self):
        self.receiver.close()
        self.close()
        self.connected = False
        self.trace(self.IN, "Connection closed")

channel_counter = 0

class proxy_receiver(xapian_message_pipe, asynchat.async_chat):
    ac_in_buffer_size = 32768
    ac_out_buffer_size = 32768
    def __init__ (self, server, (conn, addr)):
        asynchat.async_chat.__init__(self, conn)
        xapian_message_pipe.__init__(self)
        self.server = server
        global channel_counter
        self.id = channel_counter
        channel_counter += 1
        self.sender = proxy_sender (self, server.there, server.options)
        self.sender.id = self.id
        self.npackets = 0 # number of *complete* requests.
        self.trace(self.IN, "Connection from %s", addr)

    def collect_incoming_data (self, data):
        msg_type, nbytes = self.handle_incoming_data(data, self.sender)
        if msg_type is None:
            return
        self.npackets += 1
        if self.server.options.simulate_disconnect_count:
            if self.npackets % self.server.options.simulate_disconnect_count == 0:
                self.trace(self.OUT, "** simulating connection failure **")
                self.close()
                self.sender.close()
                self.connected = False
                return
        self.trace(self.IN, "%s (%d bytes)", message_types[msg_type], nbytes)

    def handle_close (self):
        self.trace(self.IN, "client disconnected")
        self.connected = False
        self.close()
        self.sender.close()

def run_proxy_server(opts):
    ps = proxy_server (opts.remote_host, opts.remote_port, opts.local_port, opts)
    print "Waiting for a connection..."
    try:
        asyncore.loop()
    except asyncore.ExitNow:
        print "proxy server terminating"

def main():
    parser = optparse.OptionParser("%prog [options]",
                                   description=__doc__)
    parser.add_option("-v", "--verbose",
                      action="count", dest="verbose", default=1,
                      help="Increase level of verbosity")

    parser.add_option("-l", "--local-port", type="int",
                      help="Local port to listen for connections on")

    parser.add_option("-p", "--remote-port", type="int",
                      help="Remote port to proxy")

    parser.add_option("-r", "--remote-host",
                      help="Remote host to proxy (default=localhost)",
                      default="localhost")

    # Debugging options: all called 'simulate' as they are attempting
    # to simulate real network conditions.
    parser.add_option("", "--simulate-disconnect-count", type="int",
                      help="""If specified, the proxy will disconnect from
                              the client after this many requests.  Useful
                              for testing auto-reconnect semantics""")

    opts, args = parser.parse_args()
    if not opts.local_port or not opts.remote_port:
        parser.error("Must specify --local-port and --remote-port")
    if args:
        parser.error("Don't specify args")

    run_proxy_server(opts)

if __name__=='__main__':
    main()
